#include "Chebyshev.hpp"
#include "par_CSRMatrix.hpp"

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
ChebyshevSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::ChebyshevSolver()
{
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void ChebyshevSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetCheckStep(int step)
{
    ck_step = step;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void ChebyshevSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetRatio(double ratio)
{
    eig_ratio = ratio;
}

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void ChebyshevSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetMaxEigenValue(double max_eig)
{
    lambda = max_eig;
}

/*
template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void ChebyshevSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::SetOperator(const Operator<idx_t, ksp_t, dof> & op)
{
    this->oper = & op;

    this->setup_time = 0.0;
    // 注意，当ksp_t与pc_t的实例化类型不同时，此处会编译报错，需要另外改写：注意是用高精度的完全setup之后再截断
    if (this->prec) this->prec->SetOperator(*(this->oper));
}
*/

template<typename idx_t, typename pc_data_t, typename pc_calc_t, typename ksp_t, int dof>
void ChebyshevSolver<idx_t, pc_data_t, pc_calc_t, ksp_t, dof>::Mult_ipl(
    const par_Vector<idx_t, ksp_t, dof> & b, par_Vector<idx_t, ksp_t, dof> & x) const
{
    assert(this->oper != nullptr);

    MPI_Comm comm = x.comm;
    int my_pid; MPI_Comm_rank(comm, &my_pid);
    if (my_pid == 0) printf("Chebyshev with check len = %d, max_eigen = %.5e, ratio = %.3e\n", ck_step, lambda, eig_ratio);
    double * record = this->part_times;
    memset(record, 0, sizeof(double) * NUM_KRYLOV_RECORD);
    par_Vector<idx_t, ksp_t, dof> r(b);
    par_Vector<idx_t, pc_calc_t, dof> * pc_buf_b = nullptr, * pc_buf_x = nullptr;
    if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
        pc_buf_b = new par_Vector<idx_t, pc_calc_t, dof>(b.comm, b.glb_nrows, b.beg_row, b.end_row);
        pc_buf_x = new par_Vector<idx_t, pc_calc_t, dof>(*pc_buf_b);
        pc_buf_b->set_val(0.0);
        pc_buf_x->set_val(0.0);
    }
    double b_norm = this->Norm(b);
    double r_norm = 0.0;

    if (this->zero_guess) {
        vec_copy(b, r);
        r_norm = b_norm;
    } else {
        this->oper->Mult(x, r, false);// r <- A*x
        vec_add(b, (ksp_t) -1.0, r, r);// r <- b - A*x
        r_norm = this->Norm(r);
    }
    if (my_pid == 0) {
        printf("L2 norm of b: %20.16e\n", (double)b_norm);
        printf("Initial L2 norm of residual: %20.16e\n", (double)r_norm);
    }

    double alpha = lambda;
    double beta  = lambda * eig_ratio;

    double theta = 2.0 / (beta + alpha);
    double delta = 2.0 / (beta - alpha);
    double sigma = delta / theta;
    double rho   = 1.0 / sigma;
    double rho1  = 0.0;

    par_Vector<idx_t, ksp_t, dof> z(b), p(b);
    int & k = this->final_iter;
    for ( k = 0; k < this->max_iter; k ++ ) {
        // 求解残差方程 M*z = r
        if (this->oper) {
            bool zero_guess = true;
            if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                assert(false);
            } else {
                if (zero_guess) p.set_val(0.0);
                record[PREC] -= wall_time();
                this->prec->Mult(r, p, zero_guess);
                record[PREC] += wall_time();
            }
        }
        else vec_copy(r, p);

        if (k == 0) {
            for (idx_t i = 0; i < (r.end_row - r.beg_row)*dof; i++)
                z.local_vector->data[i] = theta * p.local_vector->data[i];
        } else {
            rho = 1.0 / (2.0 * sigma - rho1);
            for (idx_t i = 0; i < (r.end_row - r.beg_row)*dof; i++)
                z.local_vector->data[i] = (2.0*rho*delta) * p.local_vector->data[i]
                                        + (    rho*rho1 ) * z.local_vector->data[i];
        }

        vec_add(x, (ksp_t) 1.0, z, x);
        rho1 = rho;

        // 计算残差
        this->oper->Mult(x, r, false);
        vec_add(b, (ksp_t)-1.0, r, r);

        if ( k % ck_step == 0 ) {
            r_norm = this->Norm(r);
            double rel_norm = r_norm / b_norm;
            if (my_pid == 0) printf("Iter %d ||r|| %20.16e ||r||/||b|| %20.16e\n", k, r_norm, rel_norm);
            if (r_norm < this->abs_tol || rel_norm < this->rel_tol) {
                break;
            }
        }
    }

/*
    const double max_eig = lambda;
    const double min_eig = lambda * eig_ratio;
    const double d = 0.5 * (max_eig + min_eig);
    const double c = 0.5 * (max_eig - min_eig);
    par_Vector<idx_t, ksp_t, dof> z(b), p(b);

    int & i = this->final_iter;
    double alpha = 0.0;
    for ( i = 1; i < this->max_iter; i ++) {
        // 求解残差方程 M*z = r
        if (this->oper) {
            bool zero_guess = true;
            if constexpr (sizeof(pc_calc_t) != sizeof(ksp_t)) {
                VecTransPrec(r, *pc_buf_b);// copy in
                if (zero_guess) pc_buf_x->set_val(0.0);
                else VecTransPrec(z, *pc_buf_x);// copy in
                record[PREC] -= wall_time();
                this->prec->Mult(*pc_buf_b, *pc_buf_x, zero_guess);
                record[PREC] += wall_time();
                VecTransPrec(*pc_buf_x, z);
            } else {
                if (zero_guess) z.set_val(0.0);
                record[PREC] -= wall_time();
                this->prec->Mult(r, z, zero_guess);
                record[PREC] += wall_time();
            }
        }
        else vec_copy(r, z);

        if (i == 1) {
            vec_copy(z, p);
            alpha = 1.0 / d;
        } else if (i == 2) {
            double beta = 0.5 * (c * alpha) * (c * alpha);
            alpha = 1.0 / (d - beta / alpha);
            vec_add(z, (ksp_t) beta, p, p);
        } else {
            double beta = (0.5 * c * alpha) * (0.5 * c * alpha);
            alpha = 1.0 / (d - beta / alpha);
            vec_add(z, (ksp_t) beta, p, p);
        }

        vec_add(x, (ksp_t) alpha, p, x);// x <- x + alpha*p
        record[OPER] -= wall_time();
        this->oper->Mult(x, r, false);
        record[OPER] += wall_time();
        vec_add(b, (ksp_t) -1.0, r, r);// 计算当前残差

        if ( i % ck_step == 0 ) {
            r_norm = this->Norm(r);
            double rel_norm = r_norm / b_norm;
            if (my_pid == 0) printf("Iter %d ||r|| %20.16e ||r||/||b|| %20.16e\n", i, r_norm, rel_norm);
            if (r_norm < this->abs_tol || rel_norm < this->rel_tol) {
                break;
            }
        }
    }
*/
/*
    const double max_eig = lambda;
    const double min_eig = lambda * eig_ratio;
    double beta_0 = 4.0 / (max_eig + min_eig);
    vec_add(x, (ksp_t)(0.5*beta_0), r, x_m);// x1 = x0 + 1/2*beta_0*r_0
    double beta_m = beta_0;
    par_Vector<idx_t, ksp_t, dof> x_m(b);
    int m = 1;
    for ( ; m < this->max_iter; m ++) {
        this->oper->Mult(x_m, r, false);
        vec_add(b, (ksp_t) -1.0, r, r);// 计算残差r_m

        if (m % ck_step == 0) {// 检查是否收敛
            r_norm = this->Norm(r);
            double rel_norm = r_norm / b_norm;
            if (my_pid == 0) printf("Iter %d ||r|| %20.16e ||r||/||b|| %20.16e\n", m, r_norm, rel_norm);
            if (r_norm < this->abs_tol || rel_norm < this->rel_tol) {
                break;
            }
        }

        double bm_inv = 0.5*(max_eig + min_eig) - 0.0625*(max_eig - min_eig)*(max_eig - min_eig)*beta_m;// (beta_m)^{-1}
        double bm = 1.0 / bm_inv;
        double am = 0.5*(max_eig + min_eig)*bm;// alpha_m
        double one_minus_am = 1.0 - am;// 1.0 - alpha_m

        for (idx_t i = 0; i < (r.end_row - r.beg_row) * dof; i ++) {
            ksp_t x_m_plus1 =
                    am           * x_m.local_vector->data[i]
                +   one_minus_am * x  .local_vector->data[i]
                +   bm           * r  .local_vector->data[i];
            // move 
            x  .local_vector->data[i] = x_m.local_vector->data[i];
            x_m.local_vector->data[i] = x_m_plus1;
        }
        // 更新
        beta_m = bm;
    }
    vec_copy(x_m, x);
    this->final_iter = m;
*/

    delete pc_buf_b; delete pc_buf_x;
}


template class ChebyshevSolver<int, double, double, double, 1>;
template class ChebyshevSolver<int, double, double, double, 2>;
template class ChebyshevSolver<int, double, double, double, 3>;
template class ChebyshevSolver<int, double, double, double, 4>;

template class ChebyshevSolver<int, float , double, double, 1>;
template class ChebyshevSolver<int, float , double, double, 2>;
template class ChebyshevSolver<int, float , double, double, 3>;
template class ChebyshevSolver<int, float , double, double, 4>;

template class ChebyshevSolver<int, float , float , double, 1>;
template class ChebyshevSolver<int, float , float , double, 2>;
template class ChebyshevSolver<int, float , float , double, 3>;
template class ChebyshevSolver<int, float , float , double, 4>;

template class ChebyshevSolver<int, float , float , float , 1>;
template class ChebyshevSolver<int, float , float , float , 2>;
template class ChebyshevSolver<int, float , float , float , 3>;
template class ChebyshevSolver<int, float , float , float , 4>;
#ifdef USE_FP16
template class ChebyshevSolver<int, __fp16, float , double, 1>;
template class ChebyshevSolver<int, __fp16, float , double, 2>;
template class ChebyshevSolver<int, __fp16, float , double, 3>;
template class ChebyshevSolver<int, __fp16, float , double, 4>;

template class ChebyshevSolver<int, __fp16, float , float , 1>;
template class ChebyshevSolver<int, __fp16, float , float , 2>;
template class ChebyshevSolver<int, __fp16, float , float , 3>;
template class ChebyshevSolver<int, __fp16, float , float , 4>;
#endif